#include <fstream>
#include <random>

#include <bthread/bthread.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

#include "util/hashmap/hash_map.h"
#include "util/hashmap/hash_set.h"
#include "util/algo/lock.h"

using namespace hawking::indexlib;
SpinLock spinlock;
DEFINE_int32(hashmap_test_tasks, 128, "num of threads or bthreads");

void* WriteSet(void* arg) {
    std::uniform_int_distribution<int64_t> dist(0, 1000000);
    std::random_device rd;
    std::mt19937 mt(rd());
    // uint64_t idx = 0;
    HashSet<uint64_t>* hashset =
        reinterpret_cast<HashSet<uint64_t>*>(arg);
    while (1) {
        int64_t ran_int = dist(mt);
        hashset->Insert(ran_int);
        // ++idx;
        bthread_usleep(1);  // 1 us
    }

    return nullptr;
}

void* ReadSet(void* arg) {
    std::uniform_int_distribution<int64_t> dist(0, 1000000);
    std::random_device rd;
    std::mt19937 mt(rd());
    HashSet<uint64_t>* hashset =
        reinterpret_cast<HashSet<uint64_t>*>(arg);
    uint64_t dummy = 0;
    while (1) {
        int64_t ran_int = dist(mt);
        hashset->Exist(ran_int);

        LOG_EVERY_N(INFO, 1000) << "size " << hashset->Size();
        bthread_usleep(10);  // 10 us
    }

    return nullptr;
}

TEST(hashset_test, one_write_multi_read) {
    HashSet<uint64_t> hashset(10000000);

    // write bthread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, WriteSet, &hashset);
    sleep(1);

    // read bhtreads
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_hashmap_test_tasks);
    for (size_t idx = 0; idx < FLAGS_hashmap_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, ReadSet, &hashset);
    }

    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }
}

void* Write(void* arg) {
    uint64_t idx = 0;
    HashMap<uint64_t, uint64_t>* hashmap =
        reinterpret_cast<HashMap<uint64_t, uint64_t>*>(arg);
    while (1) {
        hashmap->Insert(idx, idx);
        ++idx;
        bthread_usleep(10 * 1000);  // 10 ms
    }

    return nullptr;
}

void* Read(void* arg) {
    HashMap<uint64_t, uint64_t>* hashmap =
        reinterpret_cast<HashMap<uint64_t, uint64_t>*>(arg);
    std::unique_ptr<std::pair<uint64_t, uint64_t>> dummy =
        std::make_unique<std::pair<uint64_t, uint64_t>>();
    while (1) {
        HashMap<uint64_t, uint64_t>::Iterator iter = hashmap->CreateIterator();
        while (iter.HasNext()) {
            *dummy = iter.Next();
        }

        LOG_EVERY_N(INFO, 1000) << "size " << hashmap->Size();
    }

    return nullptr;
}

TEST(hashmap_test, one_write_multi_read) {
    HashMap<uint64_t, uint64_t> hashmap(10000000);

    // write bthread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, Write, &hashmap);
    sleep(1);

    // read bhtreads
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_hashmap_test_tasks);
    for (size_t idx = 0; idx < FLAGS_hashmap_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, Read, &hashmap);
    }

    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }
}

void* WriteComp(void* arg) {
    uint64_t idx = 0;
    std::unordered_map<uint64_t, uint64_t>* hashmap =
        reinterpret_cast<std::unordered_map<uint64_t, uint64_t>*>(arg);
    while (1) {
        {
            ScopedSpinLock mutex(spinlock);
            hashmap->emplace(idx, idx);
        }
        ++idx;
        bthread_usleep(1 * 1000);
    }

    return nullptr;
}

void* ReadComp(void* arg) {
    std::unordered_map<uint64_t, uint64_t>* hashmap =
        reinterpret_cast<std::unordered_map<uint64_t, uint64_t>*>(arg);
    std::unique_ptr<std::pair<uint64_t, uint64_t>> dummy =
        std::make_unique<std::pair<uint64_t, uint64_t>>();
    while (1) {
        {
            ScopedSpinLock mutex(spinlock);
            for (const auto& iter : *hashmap) {
                dummy->first = iter.first;
                dummy->second = iter.second;
            }
        }

        bthread_usleep(1 * 1000);
        LOG_EVERY_N(INFO, 1000) << "size " << hashmap->size();
    }

    return nullptr;
}

TEST(hashmap_test, compare_test) {
    std::unordered_map<uint64_t, uint64_t> hashmap;
    hashmap.reserve(10000000);

    // write bthread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, WriteComp, &hashmap);
    sleep(1);

    // read bhtreads
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_hashmap_test_tasks);
    for (size_t idx = 0; idx < FLAGS_hashmap_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, ReadComp, &hashmap);
    }

    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }
}

int main(int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}